train_config = {
    "all_steps": 10000,
    "save_steps": 500,
    "wandb_name":"0504",
    "wandb_project":"refine",
    "save_path":  "./ckp/GLM-4-9B-0414-refine",
    "record_path": "./reward_record.txt",
    "gen_data_path": "./gen_data.json",
    "gen_device":2,  
    "data_path":"alpaca_evol_instruct_70k_select.json",
    "beta": 0.04,
    "model_path": "Qwen2.5-7B",
    "Q_batch_size": 8,
    "num_pre_Q": 8,
    "train_batch_size":2,
    "gen_update_steps": 16,
    "compute_gen_logps": True,
    "clip_param": 0.2,
    "ref_server": "http://localhost:59864",
    "port": 59864,
    "wandb_key":"XXX",
    "eval_prompt":"""Now, I want to test an AI assistant's ability to answer questions.
Below is a question, a ground truth answer, and an answer generated by an AI assistant.
Please rate the AI assistant's answers according to the ground truth answer.
If you think the answer is correct, your output is 1; otherwise, your output is 0.
Your output is just 0 or 1.
"""
}
prompt_config = {
    "refine_prompt":'''You are a helpful assistant with self-refinement capability. 
After the user asks a question, you first think carefully and then give the answer.
The thinking process and answer should be enclosed within <think> </think> and <answer> </answer> tags respectively. Note that you can only use once these four tags.
In the <think> and </think> tag,follow these rules:
1. Start with an intial thought process on how to approach the question.
2. when you determine that additional clarification, detail, or improved reasoning is necessary, insert <refine> </refine> tag and then specify what needs to be reconsidered or improved. You can use both tags multiple times.
3. Continue to advance your reasoning after each refinement until you feel there is no more room for improvement.
This is how your full response should be structured:
<think>Here is your thinking process, when you think you need to reflect, insert <refine>your reflection</refine>.Repeat the iterative process as many times as necessary before moving to the final answer.</think><answer>Here is an answer at the end of the thinking process.</answer> ''',
    "raw_prompt": "You are a helpful AI assistant. Given a task description, please provide the corresponding response to the request."
}
eval_config = {
    "eval_llm_port":59811,
}
ds_config = {
    "train_micro_batch_size_per_gpu": train_config['train_batch_size'],
    "gradient_accumulation_steps": 4,
    "steps_per_print": 5,
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": 1e-6 }
    },
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": False,
        "reduce_scatter": True,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": True,
        "stage3_gather_16bit_weights_on_model_save": True,
        "offload_optimizer": {"device": "cpu"}
    }
}